import random
from copy import copy
import math
import time

class GA:
  ROULETTE = 'ROULETTE'
  TOURNAMENT = 'TOURNAMENT'
  
  class Param:
    BOOL = 'BOOL'
    INT = 'INT'
    FLOAT = 'FLOAT'
    ENUM = 'ENUM'
    FLOAT_LIST = 'FLOAT_LIST'
    INT_LIST = 'INT_LIST'
    
    def __init__( self, name, type, size, **kwargs ):
      self.name = name
      self.type = type
      self.size = size
      self.__dict__.update( kwargs )
      
    @staticmethod
    def bool( name ):
      return GA.Param( name, GA.Param.BOOL, 1 )
      
    @staticmethod
    def int( name, min, max ):
      return GA.Param( name, GA.Param.INT, GA.Param.bitsForInt( max - min ), min = min, max = max )
      
    @staticmethod
    def float( name, min, max, acc ):
      steps = int( math.ceil( ( max - min ) / acc ) )
      return GA.Param( name, GA.Param.FLOAT, GA.Param.bitsForInt( steps ), min = min, max = max )
      
    @staticmethod
    def enum( name, vals ):
      return GA.Param( name, GA.Param.ENUM, GA.Param.bitsForInt( len( vals ) ), vals = vals )
      
    @staticmethod
    def floatList( name, min, max, acc, num ):
      steps = int( math.ceil( ( max - min ) / acc ) )
      numBits = GA.Param.bitsForInt( steps )
      return GA.Param( name, GA.Param.FLOAT_LIST, numBits * num, min = min, max = max, num = num, numBits = numBits )
      
    @staticmethod
    def intList( name, min, max, num ):
      numBits = GA.Param.bitsForInt( max - min )
      return GA.Param( name, GA.Param.INT_LIST, numBits * num, min = min, max = max, num = num, numBits = numBits )
      
    @staticmethod
    def bitsForInt( i ):
      return int( math.ceil( math.log( i, 2 ) ) )
      
    @staticmethod
    def decodeInt( bits ):
      return sum( [ 2 ** i * b for i, b in enumerate( bits ) ] )
      
    @staticmethod
    def decodeFloat( min, max, bits ):
      return min + ( max - min ) * GA.Param.decodeInt( bits ) / float( 2 ** len( bits ) - 1 )
      
    def decode( self, bits ):
      if self.type == GA.Param.BOOL:
        return bits[0]
      elif self.type == GA.Param.INT:
        return self.min + GA.Param.decodeInt( bits ) % ( self.max - self.min + 1 )
      elif self.type == GA.Param.FLOAT:
        return GA.Param.decodeFloat( self.min, self.max, bits )
      elif self.type == GA.Param.ENUM:
        return self.vals[ GA.Param.decodeInt( bits ) % len( self.vals ) ]
      elif self.type == GA.Param.FLOAT_LIST:
        return [ GA.Param.decodeFloat( self.min, self.max, bits[ i: i + self.numBits ] ) for i in range( len( bits ) )[::self.numBits] ]
      elif self.type == GA.Param.INT_LIST:
        return [ self.min + GA.Param.decodeInt( bits[ i: i + self.numBits ] ) % ( self.max - self.min + 1 ) for i in range( len( bits ) )[::self.numBits] ]
      else:
        raise NotImplemented
  
  def __init__( self, params, eval, popSize = 1000, numGens = None, maxTime = None, crossProb = 0.25, mutateProb = 0.01, selectionMethod = None, verbose = False ):
    self.params = params
    self.genomeLength = sum( [ param.size for param in self.params ] )
    self.eval = eval
    self.popSize = popSize
    self.numGens = numGens
    self.maxTime = maxTime
    assert self.numGens or self.maxTime
    self.crossProb = crossProb
    self.mutateProb = mutateProb
    self.selectionMethod = selectionMethod or GA.TOURNAMENT
    self.verbose = verbose
    
  def decode( self, ind ):
    pos = 0
    vals = {}
    for param in self.params:
      vals[ param.name ] = param.decode( ind[ pos: pos + param.size ] )
      pos += param.size
    return vals
    
  def generate( self ):
    return [ random.randint( 0, 1 ) for i in range( self.genomeLength ) ]
    
  def combine( self, pop, fitness ):
    newPop = []
    for i in range( 0, len( pop ), 2 ):
      i1 = self.select( pop, fitness )
      i2 = self.select( pop, fitness )
      i1, i2 = self.cross( copy( i1 ), copy( i2 ) )
      self.mutate( i1 )
      self.mutate( i2 )
      newPop += [ i1, i2 ]
    return newPop
    
  def cross( self, i1, i2 ):  
    if random.random() < self.crossProb:
      pos = random.randint( 0, min( len( i1 ), len( i2 ) ) - 1 )
      return i1[:pos] + i2[pos:], i2[:pos] + i1[pos:]
    else:
      return i1, i2
      
  def mutate( self, ind ):
    for i in range( len( ind ) ):
      if random.random() < self.mutateProb:
        ind[i] = int( not ind[i] )
        
  def select( self, pop, fitness ):
    if self.selectionMethod == GA.ROULETTE:
      minFitness = min( fitness )
      if minFitness < 0:
        fitness = [ score - minFitness for score in fitness ]
      total = sum( fitness )
      target = random.uniform( 0, total )
      inc = 0
      for ind, score in zip( pop, fitness ):
        inc += score
        if inc >= target:
          return ind
      raise RuntimeError
    elif self.selectionMethod == GA.TOURNAMENT:
      i1 = random.randint( 0, len( pop ) - 1 )
      i2 = random.randint( 0, len( pop ) - 1 )
      if fitness[i1] > fitness[i2]:
        return pop[i1]
      else:
        return pop[i2]
    else:
      raise NotImplemented( self.selectionMethod )
    
  def run( self ):
    pop = [ self.generate() for i in range( self.popSize ) ]
    best = None
    
    gen = 0
    start = time.time()
    while ( self.numGens and gen < self.numGens ) or ( self.maxTime and time.time() - start < self.maxTime ):
      gen += 1
      decodedInds = map( self.decode, pop )
      fitness = map( self.eval, decodedInds )
      
      for score, bits, ind in zip( fitness, pop, decodedInds ):
        if not best or score > bestScore:
          bestScore = score
          best = copy( ind )
          bestGen = gen
          bestBits = bits
      
      if self.verbose:
        print '%d %.2f %.2f %.2f %.2f %.2f %d' % ( gen, time.time() - start, min( fitness ), sum( fitness ) / float( len( fitness ) ), max( fitness ), bestScore, bestGen )
      
      pop = self.combine( pop, fitness )

    return best
    
  @staticmethod
  def test():
    def eval( vals ):
      til = vals[ 'testIntList' ]
      return vals[ 'testBool' ] + \
             vals[ 'testInt' ] / vals[ 'testFloat' ] + \
             int( vals[ 'testEnum' ] == 5 ) + \
             -sum( [ abs( i - val ) for i, val in enumerate( vals[ 'testFloatList' ] ) ] ) + \
             int( til[0] * til[1] == til[2] * vals[ 'testInt' ] )
      
    sol = GA( [ GA.Param.bool( 'testBool' ), 
                 GA.Param.int( 'testInt', 0, 10 ),
                 GA.Param.float( 'testFloat', -20, 20, 0.1 ),
                 GA.Param.enum( 'testEnum', range( 10 ) ),
                 GA.Param.floatList( 'testFloatList', 0, 3, 0.1, 3 ),
                 GA.Param.intList( 'testIntList', 1, 10, 3 ),
               ], eval, maxTime = 5, verbose = True ).run()

    assert sol[ 'testBool' ]
    assert sol[ 'testInt' ] == 10
    assert sol[ 'testFloat' ] <= 0.1
    assert sol[ 'testEnum' ] == 5
    assert all( [ abs( val - i ) <= 0.1 for i, val in enumerate( sol[ 'testFloatList' ] ) ] )
    til = sol[ 'testIntList' ]
    assert til[0] * til[1] == til[2] * sol[ 'testInt' ]
    
if __name__ == '__main__':
  GA.test()
